import numpy as np
from matplotlib import animation
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FFMpegWriter
writer = FFMpegWriter(fps=30, metadata=dict(artist='Me'), bitrate=1800)

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# Constants.
r1 = 1.496e8            # Earth orbit radius. (km)
r2 = 2.279e8            # Mars orbit radius. (km)
R1 = 6371               # Earth radius. (km)
M = 1.989e30            # Sun mass. (kg)
M1 = 5.965e24           # Earth mass. (kg)
T2 = 686.98             # Mars cycle time. (d)
G = 6.67e-20            # Gravity constant. (kg, km, s)
a_max = 2.773e9         # The maximum of a.

fig = plt.figure()
ax = fig.gca(projection='3d')


def init():
    a1 = np.linspace(r1+1, (r1+r2)/2, num=50)
    a = np.linspace((r1+r2)/2, 5*r2, num=50)

    p1 = 2 * np.sqrt(a) / (np.sqrt(G*M))
    p2 = np.pi*a / 4
    p3 = a / 2 * np.arcsin((a-r2) / (a-a1))
    p4 = np.sqrt(2*a*r2 - 2*a*a1 + a1**2 - r2**2) / 2

    t = p1 * (p2 - p3 - p4) / (60*60*24)
    inside_arctan = np.sqrt((2*a-a1-r2) * (r2-a1) *
                            (2*a*a1-a1**2)) / (a*r2 - 2*a*a1 + a1**2)
    a1, a = np.meshgrid(a1, a)
    f = np.arctan(inside_arctan)/np.pi*180 + \
        (np.pi*np.sqrt(a1**3/(G*M))/(60*60*24) + t) / T2 * 360

    ax.set_xlabel('a1')
    ax.set_ylabel('a')
    # Plot the surface.
    surf = ax.plot_surface(a1, a, f, linewidth=0, antialiased=False)


def rotate(angle):
    ax.view_init(elev=10, azim=angle)


rotate_animation = animation.FuncAnimation(
    fig, rotate, init_func=init, frames=np.arange(0, 362, 2), interval=33)

plt.title('图像')
# plt.show()

rotate_animation.save("test1.gif", writer=writer)
